# 期权ETF Sell Put策略回测分析
# 分析sell put策略在不同时间点的收益率、胜率和赔率

import pandas as pd
import os
import sys
import glob
from datetime import datetime

# 获取脚本所在目录的父目录（项目根目录）
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(SCRIPT_DIR)

# 数据目录（使用绝对路径）
DAILY_DATA_DIR = os.path.join(PROJECT_ROOT, 'data', 'opt_etf_daily')
MIN_TRADING_DAYS = 100  # 最小交易天数
HOLDING_PERIODS = [20, 40, 60, 80]  # 持仓周期（交易日）


def load_all_daily_data(data_dir=DAILY_DATA_DIR):
    """
    加载所有期权的日线数据
    
    参数:
    data_dir (str): 数据目录路径
    
    返回:
    dict: {ts_code: DataFrame} 格式的数据字典
    """
    print(f"正在从 {data_dir} 加载数据...")
    print(f"目录是否存在: {os.path.exists(data_dir)}")
    
    if not os.path.exists(data_dir):
        print(f"错误: 数据目录不存在: {data_dir}")
        return {}
    
    all_data = {}
    
    # 获取所有CSV文件
    csv_files = glob.glob(os.path.join(data_dir, '*.csv'))
    print(f"找到 {len(csv_files)} 个数据文件")
    
    if len(csv_files) == 0:
        print("警告: 没有找到任何CSV文件")
        return {}
    
    for file_path in csv_files:
        try:
            # 从文件名提取ts_code（去掉.csv后缀，将_替换回.）
            filename = os.path.basename(file_path)
            ts_code = filename.replace('.csv', '').replace('_', '.')
            
            # 读取CSV文件
            df = pd.read_csv(file_path, encoding='utf-8-sig')
            
            if not df.empty:
                all_data[ts_code] = df
        except Exception as e:
            print(f"  读取文件 {file_path} 失败: {e}")
            import traceback
            traceback.print_exc()
    
    print(f"成功加载 {len(all_data)} 个期权的数据")
    return all_data


def filter_by_trading_days(all_data, min_days=MIN_TRADING_DAYS):
    """
    过滤掉数据行数不足指定天数的数据
    
    参数:
    all_data (dict): 所有期权数据（应该已经排序和清理过）
    min_days (int): 最小交易天数
    
    返回:
    dict: 过滤后的数据
    """
    print(f"\n正在过滤数据（最少 {min_days} 个交易日）...")
    filtered_data = {}
    filtered_count = 0
    
    for ts_code, df in all_data.items():
        if len(df) >= min_days:
            filtered_data[ts_code] = df
        else:
            filtered_count += 1
            if filtered_count <= 10:  # 只显示前10个被过滤的
                print(f"  {ts_code}: 只有 {len(df)} 个交易日，已过滤")
    
    if filtered_count > 10:
        print(f"  ... 还有 {filtered_count - 10} 个期权被过滤")
    
    print(f"过滤后剩余 {len(filtered_data)} 个期权（共过滤 {filtered_count} 个）")
    return filtered_data


def sort_by_trade_date(all_data):
    """
    对每个期权的数据按交易日期正序排序（最早的在前）
    
    参数:
    all_data (dict): 所有期权数据
    
    返回:
    dict: 排序后的数据
    """
    print("\n正在按交易日期排序...")
    sorted_data = {}
    
    for ts_code, df in all_data.items():
        # 创建副本，避免修改原数据
        df_copy = df.copy()
        
        # 确保trade_date是字符串类型
        df_copy['trade_date'] = df_copy['trade_date'].astype(str)
        
        # 按交易日期正序排序（最早的日期在前，ascending=True）
        df_sorted = df_copy.sort_values('trade_date', ascending=True).reset_index(drop=True)
        
        # 移除close为NaN或<=0的行
        df_sorted = df_sorted[df_sorted['close'].notna() & (df_sorted['close'] > 0)]
        df_sorted = df_sorted.reset_index(drop=True)
        
        sorted_data[ts_code] = df_sorted
    
    print(f"已完成 {len(sorted_data)} 个期权的排序")
    return sorted_data


def calculate_returns(df, holding_periods=HOLDING_PERIODS):
    """
    计算期权的收益率
    
    参数:
    df (DataFrame): 期权日线数据（已排序）
    holding_periods (list): 持仓周期列表
    
    返回:
    dict: {holding_period: return_rate} 格式的收益率字典
    """
    if df.empty:
        return {}
    
    # 获取第一个交易日的收盘价（建仓价格）
    first_row = df.iloc[0]
    entry_price = first_row['close']
    entry_date = first_row['trade_date']
    
    if pd.isna(entry_price) or entry_price <= 0:
        return {}
    
    returns = {}
    
    # 计算每个持仓周期的收益率
    for period in holding_periods:
        # 检查是否有足够的数据
        if len(df) > period:
            exit_row = df.iloc[period]
            exit_price = exit_row['close']
            exit_date = exit_row['trade_date']
            
            if pd.isna(exit_price) or exit_price < 0:
                # 如果平仓价格为NaN或负数，跳过
                returns[period] = None
            else:
                # Sell Put: 建仓时卖出期权（收取权利金），平仓时买入期权（支付权利金）
                # 收益率 = (建仓价格 - 平仓价格) / 建仓价格
                # 如果平仓价格更低，则盈利（正收益率）
                return_rate = (entry_price - exit_price) / entry_price
                returns[period] = {
                    'return_rate': return_rate,
                    'entry_price': entry_price,
                    'exit_price': exit_price,
                    'entry_date': entry_date,
                    'exit_date': exit_date
                }
        else:
            # 数据不足，无法计算
            returns[period] = None
    
    return returns


def analyze_all_options(all_data, holding_periods=HOLDING_PERIODS):
    """
    分析所有期权的收益率
    
    参数:
    all_data (dict): 所有期权数据（已过滤和排序）
    holding_periods (list): 持仓周期列表
    
    返回:
    DataFrame: 分析结果
    """
    print("\n开始计算收益率...")
    results = []
    
    for ts_code, df in all_data.items():
        returns = calculate_returns(df, holding_periods)
        
        result_row = {
            'ts_code': ts_code,
            'total_days': len(df),
            'first_date': df.iloc[0]['trade_date'] if not df.empty else None,
            'last_date': df.iloc[-1]['trade_date'] if not df.empty else None,
        }
        
        # 添加每个持仓周期的收益率
        for period in holding_periods:
            if period in returns and returns[period] is not None:
                result_row[f'return_{period}d'] = returns[period]['return_rate']
                result_row[f'entry_price_{period}d'] = returns[period]['entry_price']
                result_row[f'exit_price_{period}d'] = returns[period]['exit_price']
                result_row[f'entry_date_{period}d'] = returns[period]['entry_date']
                result_row[f'exit_date_{period}d'] = returns[period]['exit_date']
            else:
                result_row[f'return_{period}d'] = None
                result_row[f'entry_price_{period}d'] = None
                result_row[f'exit_price_{period}d'] = None
                result_row[f'entry_date_{period}d'] = None
                result_row[f'exit_date_{period}d'] = None
        
        results.append(result_row)
    
    results_df = pd.DataFrame(results)
    print(f"完成 {len(results_df)} 个期权的收益率计算")
    
    return results_df


def calculate_win_rate_and_payoff(results_df, holding_periods=HOLDING_PERIODS):
    """
    计算所有合约在不同时间点的胜率和赔率
    
    参数:
    results_df (DataFrame): 分析结果
    holding_periods (list): 持仓周期列表
    
    返回:
    DataFrame: 胜率和赔率统计
    """
    print("\n计算胜率和赔率...")
    stats = []
    
    for period in holding_periods:
        return_col = f'return_{period}d'
        
        # 过滤出有数据的记录
        valid_returns = results_df[return_col].dropna()
        
        if len(valid_returns) == 0:
            continue
        
        # 计算胜率（盈利次数 / 总次数）
        winning_trades = (valid_returns > 0).sum()
        losing_trades = (valid_returns < 0).sum()
        total_trades = len(valid_returns)
        win_rate = winning_trades / total_trades if total_trades > 0 else 0
        
        # 计算平均收益率
        avg_return = valid_returns.mean()
        
        # 计算平均盈利和平均亏损
        avg_profit = valid_returns[valid_returns > 0].mean() if winning_trades > 0 else 0
        avg_loss = valid_returns[valid_returns < 0].mean() if losing_trades > 0 else 0
        
        # 计算赔率（平均盈利 / 平均亏损的绝对值）
        payoff_ratio = abs(avg_profit / avg_loss) if avg_loss != 0 else 0
        
        # 计算最大盈利和最大亏损
        max_profit = valid_returns.max()
        max_loss = valid_returns.min()
        
        stats.append({
            'holding_period': f'{period}交易日',
            '总交易次数': total_trades,
            '盈利次数': winning_trades,
            '亏损次数': losing_trades,
            '胜率': win_rate,
            '平均收益率': avg_return,
            '平均盈利': avg_profit,
            '平均亏损': avg_loss,
            '赔率': payoff_ratio,
            '最大盈利': max_profit,
            '最大亏损': max_loss,
        })
    
    stats_df = pd.DataFrame(stats)
    return stats_df


def main():
    """
    主函数
    """
    print("=" * 80)
    print("期权ETF Sell Put策略回测分析")
    print("=" * 80)
    
    # 1. 加载所有日线数据
    all_data = load_all_daily_data()
    
    if not all_data:
        print("没有找到数据文件，程序退出")
        return
    
    # 2. 先排序，确保数据按时间顺序排列
    sorted_data = sort_by_trade_date(all_data)
    
    # 3. 过滤掉数据行数不足100天的数据（在排序后进行）
    filtered_data = filter_by_trading_days(sorted_data, MIN_TRADING_DAYS)
    
    if not filtered_data:
        print("过滤后没有剩余数据，程序退出")
        return
    
    # 再次排序过滤后的数据（确保顺序正确）
    final_data = sort_by_trade_date(filtered_data)
    
    # 4. 计算每个期权在不同时间点的收益率
    results_df = analyze_all_options(final_data, HOLDING_PERIODS)
    
    # 保存详细结果到data目录
    output_file = os.path.join(PROJECT_ROOT, 'data', 'opt_etf_sell_put_backtest_results.csv')
    results_df.to_csv(output_file, index=False, encoding='utf-8-sig')
    print(f"\n详细分析结果已保存到: {output_file}")
    
    # 5. 计算所有合约的胜率和赔率
    stats_df = calculate_win_rate_and_payoff(results_df, HOLDING_PERIODS)
    
    # 保存统计结果到data目录
    stats_file = os.path.join(PROJECT_ROOT, 'data', 'opt_etf_sell_put_backtest_stats.csv')
    stats_df.to_csv(stats_file, index=False, encoding='utf-8-sig')
    print(f"胜率和赔率统计已保存到: {stats_file}")
    
    # 打印统计结果
    print("\n" + "=" * 80)
    print("胜率和赔率统计结果")
    print("=" * 80)
    print(stats_df.to_string(index=False))
    print("=" * 80)
    
    # 打印一些汇总信息
    print("\n汇总信息:")
    print(f"  分析的期权数量: {len(results_df)}")
    for period in HOLDING_PERIODS:
        return_col = f'return_{period}d'
        valid_count = results_df[return_col].notna().sum()
        print(f"  {period}交易日: {valid_count} 个有效数据")
    
    print("\n分析完成！")


if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(f"程序执行出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

